import numpy as np

from src.utils.task import Task
from .transformation import Transformation


class PriorityTasks:
    def __init__(self, model_nv) -> None:
        super().__init__()

        self.task_list: list[Task] = []
        self.name_list: list[str] = []
        self.idList: list[int] = []
        self.parent_id_list: list[int] = []
        self.child_id_list: list[int] = []
        self.out_delta_q = np.zeros(model_nv)
        self.out_dq = np.zeros(model_nv)
        self.out_ddq = np.zeros(model_nv)
        self.start_id = 0

    def addTask(self, name: str) -> None:
        self.task_list.append(Task(name))
        self.task_list[-1].id = len(self.task_list) - 1
        self.name_list.append(name)

    def getId(self, name: str) -> int:
        if name in self.name_list:
            return self.name_list.index(name)
        return -1

    def build_priority(self, task_order: list[str]) -> None:
        self.start_id = self.getId(task_order[0])
        for i in range(len(task_order)):
            id_cur = self.getId(task_order[i])
            if i == 0:
                self.task_list[id_cur].parent_id = -1
            else:
                id_bef = self.getId(task_order[i - 1])
                self.task_list[id_cur].parent_id = id_bef
            if i == len(task_order) - 1:
                self.task_list[id_cur].child_id = -1
            else:
                id_nxt = self.getId(task_order[i + 1])
                self.task_list[id_cur].child_id = id_nxt

    def calculate(self, des_delta_q: np.ndarray, des_dq: np.ndarray, des_ddq: np.ndarray, dyn_M: np.ndarray,
                  dyn_M_inv: np.ndarray, dq: np.ndarray):
        cur_id = self.start_id
        parent_id = self.task_list[cur_id].parent_id
        child_id = self.task_list[cur_id].child_id
        for i in range(len(self.task_list)):
            if parent_id == -1:
                self.task_list[cur_id].N = np.eye(self.task_list[cur_id].J.shape[1])
                self.task_list[cur_id].Jpre = self.task_list[cur_id].J @ self.task_list[cur_id].N
                self.task_list[cur_id].delta_q = des_delta_q + Transformation.pseudoInv_right_weighted(
                    self.task_list[cur_id].Jpre, self.task_list[cur_id].W) @ self.task_list[cur_id].errX
                self.task_list[cur_id].dq = des_dq.copy()
                ddxcmd = self.task_list[cur_id].ddxDes + self.task_list[cur_id].kp @ self.task_list[cur_id].errX + \
                         self.task_list[
                             cur_id].kd @ self.task_list[cur_id].derrX
                self.task_list[cur_id].ddq = des_ddq + Transformation.dyn_pseudoInv(self.task_list[cur_id].Jpre, dyn_M_inv,
                                                                                    True) @ (
                                                     ddxcmd - self.task_list[cur_id].dJ @ dq)
            else:
                self.task_list[cur_id].N = self.task_list[parent_id].N @ (
                        np.eye(self.task_list[parent_id].Jpre.shape[1]) - Transformation.pseudoInv_right_weighted(
                    self.task_list[parent_id].Jpre, self.task_list[parent_id].W) @ self.task_list[parent_id].Jpre)
                self.task_list[cur_id].Jpre = self.task_list[cur_id].J @ self.task_list[cur_id].N
                self.task_list[cur_id].delta_q = self.task_list[
                                                     parent_id].delta_q + Transformation.pseudoInv_right_weighted(
                    self.task_list[cur_id].Jpre, self.task_list[cur_id].W) @ (
                                                         self.task_list[cur_id].errX - self.task_list[cur_id].J @
                                                         self.task_list[parent_id].delta_q)
                self.task_list[cur_id].dq = self.task_list[parent_id].dq + Transformation.pseudoInv_right_weighted(
                    self.task_list[cur_id].Jpre, self.task_list[cur_id].W) @ (
                                                    self.task_list[cur_id].dxDes - self.task_list[cur_id].J @
                                                    self.task_list[
                                                        parent_id].dq)
                ddxcmd = self.task_list[cur_id].ddxDes + self.task_list[cur_id].kp @ self.task_list[cur_id].errX + \
                         self.task_list[
                             cur_id].kd @ self.task_list[cur_id].derrX
                self.task_list[cur_id].ddq = self.task_list[parent_id].ddq + Transformation.dyn_pseudoInv(
                    self.task_list[cur_id].Jpre, dyn_M_inv, True) @ (
                                                     ddxcmd - self.task_list[cur_id].dJ @ dq - self.task_list[
                                                 cur_id].J @
                                                     self.task_list[parent_id].ddq)

            if child_id != -1:
                parent_id = cur_id
                cur_id = child_id
                child_id = self.task_list[cur_id].child_id
            else:
                break
        self.out_delta_q[:] = self.task_list[cur_id].delta_q
        self.out_dq[:] = self.task_list[cur_id].dq
        self.out_ddq[:] = self.task_list[cur_id].ddq
        self.out_ddq[self.out_ddq > 1000000.0] = 1000000.0
        self.out_ddq[self.out_ddq < -1000000.0] = -1000000.0
